# -*- coding: UTF-8 -*-
"""
@项目名称：del_test.py
@作   者：陆地起飞全靠浪
@创建日期：2025-09-12-17:33
"""
import json
import os
import shutil
import time
import unittest
from glob import glob
import rdkit
from rdkit import Chem
import pandas as pd
from rabbit_mq.rmq_config import logger
from ChemTools.get_pocket import get_pocket
from ChemTools.utiltools import ligand_format_conversion
from ChemTools.get_chain_class import get_chain_class
from ChemTools.pdb_add_h import add_hydrogens
from pymol import cmd
from zipfile import ZipFile


class DelTest(unittest.TestCase):
    def test_del_pycache(self):
        # python3  -m unittest del_test.DelTest.test_del_pycache
        parent_dir = os.path.dirname(__file__)
        for i in range(10):
            pycache_path = os.path.join(parent_dir, '*/' * i, '__pycache__')
            os.system(f'rm -rf {pycache_path}')
            print('Delete:', pycache_path)

    def test_scaffold(self):
        from rdkit import Chem
        from rdkit.Chem import Draw
        from rdkit.Chem.Scaffolds import MurckoScaffold
        m = Chem.MolFromSmiles('O=C(NCc1cc(OC)c(O)cc1)Cc1cocc1CC')
        core = MurckoScaffold.GetScaffoldForMol(m)
        m_core = [m, core, Chem.MolFromSmiles('[*][C@H]1CC[C@@H]([C@@H](O)N2CCC[C@H](C3NC4CCCCC4N3)C2)C1')]
        pil_img = Draw.MolsToGridImage(m_core, subImgSize=(250, 250))
        for m in m_core:
            print(Chem.MolToSmiles(m))
        pil_img.show()

    def test_pdf2smi(self):
        for pdb_path in glob('/4T/AIDDserver/DataSrc/测试用例/extracted_chains_class/chain_配体_*.pdb'):
            mol = Chem.rdmolfiles.MolFromPDBFile(pdb_path)
            smi = Chem.MolToSmiles(mol)
            print(smi)

    def test_pymol(self):
        cmd.load("/4T/其他数据/智慧药物/HitDiscovery/4yzm/4yzm.pdb", "all")
        cmd.save("/4T/其他数据/智慧药物/HitDiscovery/4yzm/test_pymol/4yzm.pdb", "all")

    def test_protein_pocket_size(self):
        print('获取排序最高的蛋白质文件路径')
        save_chain_dir = "/4T/AIDDserver/MoleculeGenerate/GenerateResult/generate_2025_10_24-17_46_34/4nf4/chains"
        ligand_list = glob(os.path.join(save_chain_dir, 'chain_配体_*.pdb'))
        protein_list = glob(os.path.join(save_chain_dir, 'chain_蛋白质_*.pdb'))
        protein_pocket_size_list = []
        for protein_pdb in protein_list:
            for ligand_x in ligand_list:
                ligand_add_h_path = f"{os.path.join(save_chain_dir, os.path.basename(ligand_x)[:-4])}_ligand_addh.pdb"
                add_hydrogens(ligand_x, ligand_add_h_path)
                reference_ligand_mol2 = os.path.join(save_chain_dir, os.path.basename(ligand_x)[:-4]) + '.mol2'
                reference_ligand_sdf = os.path.join(save_chain_dir, os.path.basename(ligand_x)[:-4]) + '.sdf'
                ligand_format_conversion(ligand_add_h_path, reference_ligand_mol2)
                ligand_format_conversion(ligand_add_h_path, reference_ligand_sdf)
                pocket_pdb = os.path.join(save_chain_dir, os.path.basename(ligand_x)[:-4]) + '_pocket_ligH12A.pdb'
                get_pocket(reference_ligand_sdf, protein_pdb, reference_ligand_mol2, pocket_pdb, exwithin=12)
                protein_pocket_size_list.append([os.path.getsize(pocket_pdb), protein_pdb, ligand_x])
        _, protein_max_path, ligand_max_path = sorted(protein_pocket_size_list, key=lambda x: x[0])[-1]
        print()

    def test_smi2csv(self):
        smi_csv_df = pd.read_csv('/4T/AIDDserver/DockInput/4bvw/molecules_weight_below_300.csv')
        smi_df = pd.DataFrame(smi_csv_df.values[:3000], columns=smi_csv_df.columns.tolist())
        smi_df.to_csv(f'/4T/AIDDserver/DockInput/4bvw/ligand_unknown.csv', index=False)  # index=False不补首列序列

    def test_pandas_csv_df(self):
        import pandas as pd
        csv_df = pd.read_csv("./GradioData/davis_test.csv")  # ,header=None 首行不为标签
        city = pd.DataFrame(csv_df.values[:100], columns=csv_df.columns.tolist())
        city.to_csv(f'./GradioData/{time.strftime("%Y-%m-%d_%H-%M-%S")}_test.csv', index=False)  # index=False不补首列序列

        print()

    def test_protein_preprocessing(self):
        '''
        /4T/AIDDserver/DataSrc/lpa/4bvw.pdb
        /4T/AIDDserver/DataSrc/lpa/8tce.pdb
        /4T/AIDDserver/DataSrc/lpa/8v9b.pdb
        /4T/AIDDserver/DataSrc/lpa/8v8z.pdb
        '''
        pdb_id = '8v9b'
        ligand_save_dir = f'/4T/AIDDserver/DataSrc/lpa_test/{pdb_id}'
        print('protein格式cif转为pdb')
        input_protein_file_path = os.path.join(ligand_save_dir, f'{pdb_id}.cif')
        output_protein_file_path = os.path.join(ligand_save_dir, f'{pdb_id}.pdb')
        ligand_format_conversion(input_protein_file_path, output_protein_file_path)
        print('提取蛋白质文件中所有的chain，并分类保存到指定文件夹')
        input_protein_file_path = os.path.join(ligand_save_dir, f'{pdb_id}.pdb')  # 指定PDB文件路径（请替换为你的实际文件路径）
        save_chain_dir = os.path.join(ligand_save_dir, 'chains')  # 输出目录名称
        get_chain_class(input_protein_file_path, save_chain_dir)
        print('获取排序最高的蛋白质文件路径')
        protein_path_list = glob(os.path.join(save_chain_dir, '*_蛋白质_*.pdb'))
        protein_size_list = [[x.split('_')[-1][:-4], x] for x in protein_path_list]
        protein_max_path = sorted(protein_size_list, key=lambda x: x[0], reverse=True)[0][1]
        print('获取排序最高的配体文件路径')
        ligand_path_list = glob(os.path.join(save_chain_dir, '*_配体_*.pdb'))
        ligand_size_list = [[x.split('_')[-1][:-4], x] for x in ligand_path_list]
        ligand_max_path = sorted(ligand_size_list, key=lambda x: x[0], reverse=True)[0][1]
        print('删除旧氢，添加新氢')
        ligand_add_h_path = f"{os.path.join(ligand_save_dir, pdb_id)}_ligand_addh.pdb"
        add_hydrogens(ligand_max_path, ligand_add_h_path)
        print('复制文件到结构性目录中')
        reference_ligand_sdf = f"{os.path.join(ligand_save_dir, 'glide_pos', pdb_id)}_ligand.sdf"
        reference_ligand_mol2 = f"{os.path.join(ligand_save_dir, 'protein', pdb_id)}_ligand.mol2"
        protein_pdb = f"{os.path.join(ligand_save_dir, 'protein', pdb_id)}_protein.pdb"
        ligand_format_conversion(ligand_add_h_path, reference_ligand_sdf)
        ligand_format_conversion(ligand_add_h_path, reference_ligand_mol2)
        shutil.copy(protein_max_path, protein_pdb)
        print('提取配体口袋')
        pocket_pdb = f"{os.path.join(ligand_save_dir, 'protein', pdb_id)}_pocket_ligH12A.pdb"
        get_pocket(reference_ligand_sdf, protein_pdb, reference_ligand_mol2, pocket_pdb)
        # 配体化合物sdf、配体化合物mol2、蛋白质源文件不含配体pdb、口袋文件不含配体
        logger.info(f'预处理已完成，保存地址为：{ligand_save_dir}')
        print(reference_ligand_sdf, reference_ligand_mol2, protein_pdb, pocket_pdb)

    def test_get_pocket(self):
        for pdb_name in ['4bvw', '8tce', '8v9b', '8V8Z']:
            reference_ligand_sdf = f'/4T/AIDDserver/DataSrc/lpa/{pdb_name}/glide_pos/{pdb_name}_ligand.sdf'
            protein_pdb = f'/4T/AIDDserver/DataSrc/lpa/{pdb_name}/protein/{pdb_name}_protein.pdb'
            reference_ligand_mol2 = f'/4T/AIDDserver/DataSrc/lpa/{pdb_name}/protein/{pdb_name}_ligand.mol2'
            pocket_pdb = f'/4T/AIDDserver/DataSrc/lpa/{pdb_name}/protein/{pdb_name}_pocket_ligH12A.pdb'
            get_pocket(reference_ligand_sdf, protein_pdb, reference_ligand_mol2, pocket_pdb)

    def test_ligand_format_conversion(self):
        input_file_path = f'/4T/AIDDserver/DataSrc/lpa_test/4bvw/4bvw.cif'
        output_file_path = f'/4T/AIDDserver/DataSrc/lpa_test/4bvw/4bvw.pdb'

        ligand_format_conversion(input_file_path, output_file_path)

    def test_json(self):
        json_dumps = json.dumps({'xx': 6, 'yy': "对对对"}, ensure_ascii=False)
        print(json_dumps)
        print(json.loads(json_dumps))

    def test_xxx(self):
        print('复制文件到结构性目录中')


class GetResult(unittest.TestCase):
    def test_get_sg_smiles(self, model_SG='/4T/AIDDservle_result_-1.txt'):
        with open(model_SG, encoding='utf-8') as read_f:
            lines = read_f.read().splitlines()
        smiles_list = []
        for line in lines:
            line_dict = json.loads(line)
            for k, v in line_dict.items():
                smiles_list += v['sample']
        return smiles_list

    def test_get_pg_smiles(self, model_PG='/子生成/model_PG_generate_2025_10_24-17_51_41_SMILES_all.txt'):

        with open(model_PG, encoding='utf-8') as read_f:
            smiles_list = read_f.read().splitlines()
        return smiles_list

    def test_get_TamGen_smiles(self, model_PG_vae='4nf4_vae_flatten.tsv', model_PG_nonvae='结果/分子生成/4nf4_nonvae_flatten.tsv'):

        csv_df = pd.read_csv(model_PG_vae, delimiter='\t')  # ,header=None 首行不为标签,tsv:
        vae_list = csv_df.values[:, 0].tolist()

        csv_df = pd.read_csv(model_PG_nonvae, delimiter='\t')  # ,header=None 首行不为标签,tsv:
        nonvae_list = csv_df.values[:, 0].tolist()
        return vae_list + nonvae_list

    def test_get_MC_ado_sdf2smi(self, sdf_glob_path='/4esult/*.sdf'):
        smiles_list = []
        error_len = 0
        for sdf_path in glob(sdf_glob_path):
            try:
                mol = Chem.MolFromMolFile(sdf_path)
                smi = Chem.MolToSmiles(mol)
                smiles_list.append(smi)
            except:
                print(os.path.getsize(sdf_path), sdf_path)
                error_len += 1
        print('=' * 10, error_len)
        return smiles_list

    def test_get_all_smiles(self):
        parent_dir = '/4T/AIDDserver/非项目/Test/贵州/4NF4/结果20251028/分子生成'
        zip_file_path = f'{parent_dir}/model_MC_sample_result.zip'
        output_dir = zip_file_path[:-4]
        with ZipFile(zip_file_path, 'r') as zip_ref:
            zip_ref.extractall(output_dir)
        zip_file_path = f'{parent_dir}/model_cqu_ado_generate_result.zip'
        output_dir = zip_file_path[:-4]
        with ZipFile(zip_file_path, 'r') as zip_ref:
            zip_ref.extractall(output_dir)
        tam_gen_smiles_list = self.test_get_TamGen_smiles(model_PG_vae=f'{parent_dir}/4nf4_vae_flatten.tsv', model_PG_nonvae=f'{parent_dir}/4nf4_nonvae_flatten.tsv')
        phore_gen_smiles_list = self.test_get_pg_smiles(model_PG=f'{parent_dir}/model_PG_generate_2025_10_28-09_14_48_SMILES_all.txt')
        scaffold_GVAE_smiles_list = self.test_get_sg_smiles(model_SG=f'{parent_dir}/model_SG_sample_result_-1.txt')
        MolCRAFT_smiles_list = self.test_get_MC_ado_sdf2smi(sdf_glob_path=f'{parent_dir}/model_MC_sample_result/*.sdf')
        ado_smiles_list = self.test_get_MC_ado_sdf2smi(sdf_glob_path=f'{parent_dir}/model_cqu_ado_generate_result/*.sdf')
        all_smiles_list = tam_gen_smiles_list + phore_gen_smiles_list + scaffold_GVAE_smiles_list + MolCRAFT_smiles_list + ado_smiles_list
        all_smiles_filter = list(set(filter(lambda x: "*" not in x, all_smiles_list)))
        city = pd.DataFrame(all_smiles_filter, columns=['SMILES'])
        city.to_csv(f'{parent_dir}/generate_merge_smiles_filter.csv', index=False)  # index=False不补首列序列 tsv: delimiter='\t'


class TestEDU(unittest.TestCase):
    def test_edu_cqu_ado(self):
        import requests
        t1 = time.time()
        url = f'http://192.168.30.55:51000/run'

        '''
        <你的工程根>
        ├─ data/
        │ └─ examples_exper_c_2_19/
        ├─ example_data/
        │ └─ exper_2_19/
        │  ├─ scaffold_c.smi
        │  ├─ pocket_10A.pdb
        │  └─ scaffold_good_c.sdf
        └─ samples_exper_c_2_19/
          └─ run_001/
        '''
        task_id = 'generate_2025_10_22-03_05_24'
        parent_dir = f'/data/ltl_zhyw/code/AIDDserver/MoleculeGenerate/GenerateResult/{task_id}/ResultEduCquAdo'
        parent_dir = f'/GenerateResult/{task_id}/ResultEduCquAdo'

        examples_exper_c_2_19 = f"{parent_dir}/data/examples_exper_c_2_19"
        samples_dir = f"{parent_dir}/samples_exper_c_2_19/run_001"
        exper_2_19 = f"{parent_dir}/example_data/exper_2_19"
        for temp_path in [examples_exper_c_2_19, samples_dir, exper_2_19]:
            # os.makedirs(temp_path, exist_ok=True)
            print(f'mkdir -p {temp_path}')
        data = {
            "data_dir": examples_exper_c_2_19,  # 储存数据后处理结果的的文件夹。
            "samples_dir": samples_dir,  # 储存生成分子的文件夹。
            "scaffold_smiles_file": f"{exper_2_19}/scaffold_c.smi",  # 人工确认生成*锚点,储存需要修饰分子核心骨架smiles的文件路径。
            "protein_file": f"{exper_2_19}/pocket_10A.pdb",  # 储存靶点文件的文件路径,靶点文件必须是pdb格式。
            "scaffold_file": f"{exper_2_19}/scaffold_good_c.sdf",  # 储存需要修饰分子核心骨架立体结构的文件路径，这个文件是sdf格式的。
            "task_name": "exp",  # 任务类型直接设置"exp"就好
            "n_samples": 2  # 生成的分子数
        }
        headers = {
            'Content-Type': 'application/json',
        }
        json_data = json.dumps(data)
        sess = requests.Session()
        res = sess.post(url=url, data=json_data, headers=headers)

        if res.status_code == 200:
            result = json.loads(res.content)
            print(result, time.time() - t1)
            result = {
                'generated_sdfs': [  # 代表生成的分子路径
                    '/app/samples_exper_c_2_19/run_001/try/0/0_.sdf',
                    '/app/samples_exper_c_2_19/run_001/try/0/1_.sdf'
                ],
                'prepare': {
                    'mol_sdf': '/app/data/examples_exper_c_2_19/exp_test_full_mol.sdf',
                    'pockets_pkl': '/app/data/examples_exper_c_2_19/exp_test_full_pockets.pkl',
                    'rgroup_sdf': '/app/data/examples_exper_c_2_19/exp_test_full_rgroup.sdf',
                    'scaf_sdf': '/app/data/examples_exper_c_2_19/exp_test_full_scaf.sdf',
                    'table_csv': '/app/data/examples_exper_c_2_19/exp_test_full_table.csv'
                }
            }
        else:
            print('ERRPR status_code：', res.status_code)
